1use std::{
2 sync::{
3 Arc,
4 atomic::{AtomicBool, AtomicU32, Ordering},
5 },
6 time::Duration,
7};
8
9use crate::pp_log::{PpLog, pp_info};
10use ffmpeg_next as ffmpeg;
11use thiserror::Error as ThisError;
12
13use crate::{
14 buffer::MediaBuffer,
15 control::ControlMsg,
16 element::{Element, ElementType, Sink, Source, element_pp_log},
17 error::Result,
18 pad::SrcPad,
19};
20
21const NANOS_PER_SECOND: u128 = 1_000_000_000;
22
23#[derive(Debug, Clone, Copy)]
25pub struct AudioVolumeOptions {
26 pub gain: f32,
29 pub muted: bool,
32 pub ramp_duration: Duration,
37}
38
39impl Default for AudioVolumeOptions {
40 fn default() -> Self {
41 Self {
42 gain: 1.0,
43 muted: false,
44 ramp_duration: Duration::from_millis(10),
45 }
46 }
47}
48
49#[derive(Debug, ThisError, PartialEq)]
51pub enum AudioVolumeError {
52 #[error("gain must be finite and non-negative, got {0}")]
53 InvalidGain(f32),
54
55 #[error("gain in dB must be finite or negative infinity, got {0}")]
56 InvalidGainDb(f32),
57
58 #[error("audio frame has no channel layout")]
59 MissingChannels,
60
61 #[error("unsupported audio sample format: {0:?}")]
62 UnsupportedSampleFormat(ffmpeg::format::Sample),
63
64 #[error(
65 "audio plane {plane} is too small for its declared format: need {required} bytes, got {actual}"
66 )]
67 InvalidPlaneSize {
68 plane: usize,
69 required: usize,
70 actual: usize,
71 },
72
73 #[error(
74 "AudioVolume only processes decoded Audio frames, got a {0}; link it after an audio decoder or source"
75 )]
76 UnsupportedBuffer(&'static str),
77}
78
79#[derive(Debug)]
80struct VolumeControl {
81 gain_bits: AtomicU32,
82 muted: AtomicBool,
83}
84
85impl VolumeControl {
86 fn new(gain: f32, muted: bool) -> Self {
87 Self {
88 gain_bits: AtomicU32::new(gain.to_bits()),
89 muted: AtomicBool::new(muted),
90 }
91 }
92
93 fn gain(&self) -> f32 {
94 f32::from_bits(self.gain_bits.load(Ordering::Acquire))
95 }
96
97 fn effective_gain(&self) -> f32 {
98 if self.muted.load(Ordering::Acquire) {
99 0.0
100 } else {
101 self.gain()
102 }
103 }
104}
105
106#[derive(Debug, Clone)]
111pub struct AudioVolumeHandle {
112 control: Arc<VolumeControl>,
113}
114
115impl AudioVolumeHandle {
116 pub fn set_gain(&self, gain: f32) -> std::result::Result<(), AudioVolumeError> {
119 validate_gain(gain)?;
120 self.control
121 .gain_bits
122 .store(gain.to_bits(), Ordering::Release);
123 Ok(())
124 }
125
126 pub fn set_gain_db(&self, gain_db: f32) -> std::result::Result<(), AudioVolumeError> {
129 if gain_db == f32::NEG_INFINITY {
130 return self.set_gain(0.0);
131 }
132 if !gain_db.is_finite() {
133 return Err(AudioVolumeError::InvalidGainDb(gain_db));
134 }
135 let gain = 10.0_f32.powf(gain_db / 20.0);
136 self.set_gain(gain)
137 .map_err(|_| AudioVolumeError::InvalidGainDb(gain_db))
138 }
139
140 pub fn set_muted(&self, muted: bool) {
142 self.control.muted.store(muted, Ordering::Release);
143 }
144
145 pub fn gain(&self) -> f32 {
146 self.control.gain()
147 }
148
149 pub fn gain_db(&self) -> f32 {
150 let gain = self.gain();
151 if gain == 0.0 {
152 f32::NEG_INFINITY
153 } else {
154 20.0 * gain.log10()
155 }
156 }
157
158 pub fn is_muted(&self) -> bool {
159 self.control.muted.load(Ordering::Acquire)
160 }
161}
162
163pub struct AudioVolume {
175 pp_log: PpLog,
176 name: Arc<str>,
177 control: Arc<VolumeControl>,
178 ramp_duration: Duration,
179 current_gain: f32,
180 ramp_target: f32,
181 ramp_step: f32,
182 ramp_remaining: usize,
183 gain_envelope: Vec<f32>,
184 pad: SrcPad,
185}
186
187impl AudioVolume {
188 pub fn new(name: impl Into<String>) -> (Self, AudioVolumeHandle) {
190 Self::with_options(name, AudioVolumeOptions::default())
191 .expect("the default AudioVolumeOptions are valid")
192 }
193
194 pub fn with_options(
195 name: impl Into<String>,
196 options: AudioVolumeOptions,
197 ) -> std::result::Result<(Self, AudioVolumeHandle), AudioVolumeError> {
198 validate_gain(options.gain)?;
199 let name: Arc<str> = name.into().into();
200 let pp_log = element_pp_log(ElementType::AudioVolume, &name, None);
201 let control = Arc::new(VolumeControl::new(options.gain, options.muted));
202 let initial_gain = control.effective_gain();
203 let handle = AudioVolumeHandle {
204 control: control.clone(),
205 };
206 pp_info!(
207 pp_log: &pp_log,
208 "created: gain={}, muted={}, ramp={:?}",
209 options.gain,
210 options.muted,
211 options.ramp_duration
212 );
213 Ok((
214 Self {
215 name: name.clone(),
216 pp_log,
217 control,
218 ramp_duration: options.ramp_duration,
219 current_gain: initial_gain,
220 ramp_target: initial_gain,
221 ramp_step: 0.0,
222 ramp_remaining: 0,
223 gain_envelope: Vec::new(),
224 pad: SrcPad::new(format!("{name}_src")),
225 },
226 handle,
227 ))
228 }
229
230 fn ramp_samples(&self, sample_rate: u32) -> usize {
231 self.ramp_duration
232 .as_nanos()
233 .saturating_mul(u128::from(sample_rate))
234 .div_ceil(NANOS_PER_SECOND)
235 .min(usize::MAX as u128) as usize
236 }
237
238 fn prepare_gain_envelope(&mut self, sample_rate: u32, samples: usize) {
239 let target = self.control.effective_gain();
240 if target.to_bits() != self.ramp_target.to_bits() {
241 self.ramp_target = target;
242 self.ramp_remaining = self.ramp_samples(sample_rate);
243 if self.ramp_remaining == 0 {
244 self.current_gain = target;
245 self.ramp_step = 0.0;
246 } else {
247 self.ramp_step = (target - self.current_gain) / self.ramp_remaining as f32;
248 }
249 }
250
251 self.gain_envelope.clear();
252 self.gain_envelope.reserve(samples);
253 for _ in 0..samples {
254 if self.ramp_remaining > 0 {
255 self.current_gain += self.ramp_step;
256 self.ramp_remaining -= 1;
257 if self.ramp_remaining == 0 {
258 self.current_gain = self.ramp_target;
259 }
260 }
261 self.gain_envelope.push(self.current_gain);
262 }
263 }
264
265 fn snap_to_target(&mut self) {
266 let target = self.control.effective_gain();
267 self.current_gain = target;
268 self.ramp_target = target;
269 self.ramp_step = 0.0;
270 self.ramp_remaining = 0;
271 self.gain_envelope.clear();
272 }
273
274 fn process_audio(
275 &mut self,
276 frame: Arc<ffmpeg::frame::Audio>,
277 ) -> std::result::Result<Arc<ffmpeg::frame::Audio>, AudioVolumeError> {
278 let format = frame.format();
279 if format == ffmpeg::format::Sample::None {
280 return Err(AudioVolumeError::UnsupportedSampleFormat(format));
281 }
282 let channels = usize::from(frame.channels());
283 if channels == 0 {
284 return Err(AudioVolumeError::MissingChannels);
285 }
286
287 self.prepare_gain_envelope(frame.rate(), frame.samples());
288 if self.gain_envelope.iter().all(|gain| *gain == 1.0) {
289 return Ok(frame);
290 }
291
292 let mut frame = match Arc::try_unwrap(frame) {
296 Ok(frame) => frame,
297 Err(frame) => frame.as_ref().clone(),
298 };
299 apply_gain(&mut frame, &self.gain_envelope)?;
300 Ok(Arc::new(frame))
301 }
302}
303
304impl Element for AudioVolume {
305 fn name(&self) -> Arc<str> {
306 self.name.clone()
307 }
308
309 fn element_type(&self) -> ElementType {
310 ElementType::AudioVolume
311 }
312
313 fn pp_log(&self) -> &PpLog {
314 &self.pp_log
315 }
316
317 fn pp_log_mut(&mut self) -> &mut PpLog {
318 &mut self.pp_log
319 }
320}
321
322impl Source for AudioVolume {
323 fn src_pads(&mut self) -> &mut [SrcPad] {
324 std::slice::from_mut(&mut self.pad)
325 }
326}
327
328impl Sink for AudioVolume {
329 fn consume(&mut self, buf: MediaBuffer) -> Result<()> {
330 match buf {
331 MediaBuffer::Audio(frame) => {
332 let frame = self.process_audio(frame)?;
333 self.pad.push(MediaBuffer::Audio(frame))
334 }
335 MediaBuffer::Eos => self.pad.push(MediaBuffer::Eos),
336 MediaBuffer::Packet(_) => Err(AudioVolumeError::UnsupportedBuffer("Packet").into()),
337 MediaBuffer::Video(_) => Err(AudioVolumeError::UnsupportedBuffer("Video").into()),
338 }
339 }
340
341 fn control(&mut self, msg: ControlMsg) -> Result<()> {
342 if matches!(msg, ControlMsg::Seek(_) | ControlMsg::Stop) {
343 self.snap_to_target();
344 }
345 self.pad.control(msg)
346 }
347}
348
349fn validate_gain(gain: f32) -> std::result::Result<(), AudioVolumeError> {
350 if gain.is_finite() && gain >= 0.0 {
351 Ok(())
352 } else {
353 Err(AudioVolumeError::InvalidGain(gain))
354 }
355}
356
357fn apply_gain(
358 frame: &mut ffmpeg::frame::Audio,
359 gains: &[f32],
360) -> std::result::Result<(), AudioVolumeError> {
361 let format = frame.format();
362 let channels = usize::from(frame.channels());
363 let (planes, channels_per_plane) = if format.is_planar() {
364 (channels, 1)
365 } else {
366 (1, channels)
367 };
368 let scalar_count = gains.len().saturating_mul(channels_per_plane);
369 let required = scalar_count.saturating_mul(format.bytes());
370
371 for plane in 0..planes {
372 if plane >= frame.planes() {
381 return Err(AudioVolumeError::InvalidPlaneSize {
382 plane,
383 required,
384 actual: 0,
385 });
386 }
387 let data =
392 unsafe { std::slice::from_raw_parts_mut((*frame.as_mut_ptr()).data[plane], required) };
393 match format {
394 ffmpeg::format::Sample::U8(_) => scale_u8(data, gains, channels_per_plane),
395 ffmpeg::format::Sample::I16(_) => scale_i16(data, gains, channels_per_plane),
396 ffmpeg::format::Sample::I32(_) => scale_i32(data, gains, channels_per_plane),
397 ffmpeg::format::Sample::I64(_) => scale_i64(data, gains, channels_per_plane),
398 ffmpeg::format::Sample::F32(_) => scale_f32(data, gains, channels_per_plane),
399 ffmpeg::format::Sample::F64(_) => scale_f64(data, gains, channels_per_plane),
400 ffmpeg::format::Sample::None => {
401 return Err(AudioVolumeError::UnsupportedSampleFormat(format));
402 }
403 }
404 }
405 Ok(())
406}
407
408fn scale_u8(data: &mut [u8], gains: &[f32], channels: usize) {
409 for (index, sample) in data.iter_mut().enumerate() {
410 let centered = f32::from(*sample) - 128.0;
411 *sample = (centered * gains[index / channels] + 128.0)
412 .round()
413 .clamp(0.0, 255.0) as u8;
414 }
415}
416
417macro_rules! scale_integer {
418 ($name:ident, $sample_type:ty, $bytes:literal) => {
419 fn $name(data: &mut [u8], gains: &[f32], channels: usize) {
420 for (index, sample) in data.chunks_exact_mut($bytes).enumerate() {
421 let value = <$sample_type>::from_ne_bytes(sample.try_into().expect("chunk size"));
422 let scaled = ((value as f64) * f64::from(gains[index / channels]))
423 .round()
424 .clamp(<$sample_type>::MIN as f64, <$sample_type>::MAX as f64)
425 as $sample_type;
426 sample.copy_from_slice(&scaled.to_ne_bytes());
427 }
428 }
429 };
430}
431
432scale_integer!(scale_i16, i16, 2);
433scale_integer!(scale_i32, i32, 4);
434scale_integer!(scale_i64, i64, 8);
435
436fn scale_f32(data: &mut [u8], gains: &[f32], channels: usize) {
437 for (index, sample) in data.chunks_exact_mut(4).enumerate() {
438 let value = f32::from_ne_bytes(sample.try_into().expect("chunk size"));
439 sample.copy_from_slice(&(value * gains[index / channels]).to_ne_bytes());
440 }
441}
442
443fn scale_f64(data: &mut [u8], gains: &[f32], channels: usize) {
444 for (index, sample) in data.chunks_exact_mut(8).enumerate() {
445 let value = f64::from_ne_bytes(sample.try_into().expect("chunk size"));
446 sample.copy_from_slice(&(value * f64::from(gains[index / channels])).to_ne_bytes());
447 }
448}
449
450#[cfg(test)]
451mod tests {
452 use std::sync::Mutex;
453
454 use ffmpeg::format::sample::Type;
455
456 use super::*;
457
458 struct CapturingSink {
459 pp_log: PpLog,
460 received: Arc<Mutex<Vec<MediaBuffer>>>,
461 }
462
463 impl Element for CapturingSink {
464 fn name(&self) -> Arc<str> {
465 "capture".into()
466 }
467
468 fn element_type(&self) -> ElementType {
469 ElementType::Other
470 }
471
472 fn pp_log(&self) -> &PpLog {
473 &self.pp_log
474 }
475
476 fn pp_log_mut(&mut self) -> &mut PpLog {
477 &mut self.pp_log
478 }
479 }
480
481 impl Sink for CapturingSink {
482 fn consume(&mut self, buf: MediaBuffer) -> Result<()> {
483 self.received.lock().unwrap().push(buf);
484 Ok(())
485 }
486
487 fn control(&mut self, _msg: ControlMsg) -> Result<()> {
488 Ok(())
489 }
490 }
491
492 fn new_volume(
493 options: AudioVolumeOptions,
494 ) -> (AudioVolume, AudioVolumeHandle, Arc<Mutex<Vec<MediaBuffer>>>) {
495 let (mut volume, handle) = AudioVolume::with_options("volume", options).unwrap();
496 let received = Arc::new(Mutex::new(Vec::new()));
497 volume.src_pads()[0].link(Box::new(CapturingSink {
498 received: received.clone(),
499 pp_log: element_pp_log(ElementType::Other, "capture", None),
500 }));
501 (volume, handle, received)
502 }
503
504 fn f32_packed_frame(values: &[f32], rate: u32, channels: u16) -> Arc<ffmpeg::frame::Audio> {
505 assert_eq!(values.len() % usize::from(channels), 0);
506 let samples = values.len() / usize::from(channels);
507 let mut frame = ffmpeg::frame::Audio::new(
508 ffmpeg::format::Sample::F32(Type::Packed),
509 samples,
510 ffmpeg::ChannelLayout::default(i32::from(channels)),
511 );
512 frame.set_rate(rate);
513 frame.set_pts(Some(123));
514 let bytes = unsafe {
515 std::slice::from_raw_parts(values.as_ptr().cast::<u8>(), std::mem::size_of_val(values))
516 };
517 frame.data_mut(0)[..bytes.len()].copy_from_slice(bytes);
518 Arc::new(frame)
519 }
520
521 fn f32_planar_frame(channels_data: &[&[f32]], rate: u32) -> Arc<ffmpeg::frame::Audio> {
525 let channels = channels_data.len();
526 let samples = channels_data[0].len();
527 assert!(channels_data.iter().all(|c| c.len() == samples));
528 let mut frame = ffmpeg::frame::Audio::new(
529 ffmpeg::format::Sample::F32(Type::Planar),
530 samples,
531 ffmpeg::ChannelLayout::default(channels as i32),
532 );
533 frame.set_rate(rate);
534 frame.set_pts(Some(123));
535 for (index, values) in channels_data.iter().enumerate() {
536 frame.plane_mut::<f32>(index).copy_from_slice(values);
537 }
538 Arc::new(frame)
539 }
540
541 fn captured_f32(received: &Arc<Mutex<Vec<MediaBuffer>>>, index: usize) -> Vec<f32> {
542 let received = received.lock().unwrap();
543 let MediaBuffer::Audio(frame) = &received[index] else {
544 panic!("expected audio")
545 };
546 let count = frame.samples() * usize::from(frame.channels());
547 frame.data(0)[..count * 4]
548 .chunks_exact(4)
549 .map(|bytes| f32::from_ne_bytes(bytes.try_into().unwrap()))
550 .collect()
551 }
552
553 #[test]
554 fn mute_and_unmute_ramp_over_the_configured_number_of_samples() {
555 let options = AudioVolumeOptions {
556 ramp_duration: Duration::from_millis(10),
557 ..AudioVolumeOptions::default()
558 };
559 let (mut volume, handle, received) = new_volume(options);
560 handle.set_muted(true);
561 volume
562 .consume(MediaBuffer::Audio(f32_packed_frame(&[1.0; 10], 1_000, 1)))
563 .unwrap();
564 handle.set_muted(false);
565 volume
566 .consume(MediaBuffer::Audio(f32_packed_frame(&[1.0; 10], 1_000, 1)))
567 .unwrap();
568
569 let fade_out = captured_f32(&received, 0);
570 let fade_in = captured_f32(&received, 1);
571 for (index, sample) in fade_out.iter().enumerate() {
572 let expected = 0.9 - index as f32 * 0.1;
573 assert!((sample - expected).abs() < 1e-5, "{sample} != {expected}");
574 }
575 for (index, sample) in fade_in.iter().enumerate() {
576 let expected = 0.1 + index as f32 * 0.1;
577 assert!((sample - expected).abs() < 1e-5, "{sample} != {expected}");
578 }
579 }
580
581 #[test]
582 fn gain_db_changes_runtime_amplitude_without_changing_frame_metadata() {
583 let options = AudioVolumeOptions {
584 ramp_duration: Duration::ZERO,
585 ..AudioVolumeOptions::default()
586 };
587 let (mut volume, handle, received) = new_volume(options);
588 handle.set_gain_db(-6.0).unwrap();
589 volume
590 .consume(MediaBuffer::Audio(f32_packed_frame(
591 &[1.0, -1.0, 0.5, -0.5],
592 48_000,
593 2,
594 )))
595 .unwrap();
596
597 let gain = 10.0_f32.powf(-6.0 / 20.0);
598 let output = captured_f32(&received, 0);
599 for (actual, input) in output.iter().zip([1.0, -1.0, 0.5, -0.5]) {
600 assert!((actual - input * gain).abs() < 1e-6);
601 }
602 let received = received.lock().unwrap();
603 let MediaBuffer::Audio(frame) = &received[0] else {
604 panic!("expected audio")
605 };
606 assert_eq!(frame.rate(), 48_000);
607 assert_eq!(frame.channels(), 2);
608 assert_eq!(frame.samples(), 2);
609 assert_eq!(frame.pts(), Some(123));
610 }
611
612 #[test]
618 fn gain_scales_every_channel_of_a_planar_frame() {
619 let options = AudioVolumeOptions {
620 gain: 0.5,
621 ramp_duration: Duration::ZERO,
622 ..AudioVolumeOptions::default()
623 };
624 let (mut volume, _handle, received) = new_volume(options);
625 volume
626 .consume(MediaBuffer::Audio(f32_planar_frame(
627 &[&[1.0, -1.0], &[0.5, -0.5]],
628 48_000,
629 )))
630 .unwrap();
631
632 let received = received.lock().unwrap();
633 let MediaBuffer::Audio(frame) = &received[0] else {
634 panic!("expected audio")
635 };
636 assert_eq!(frame.plane::<f32>(0).to_vec(), vec![0.5, -0.5]);
637 assert_eq!(frame.plane::<f32>(1).to_vec(), vec![0.25, -0.25]);
638 }
639
640 #[test]
641 fn a_shared_input_is_copied_before_another_tee_branch_can_be_modified() {
642 let options = AudioVolumeOptions {
643 gain: 0.5,
644 ramp_duration: Duration::ZERO,
645 ..AudioVolumeOptions::default()
646 };
647 let (mut volume, _, received) = new_volume(options);
648 let original = f32_packed_frame(&[1.0], 48_000, 1);
649 volume
650 .consume(MediaBuffer::Audio(original.clone()))
651 .unwrap();
652
653 assert_eq!(original.plane::<f32>(0)[0], 1.0);
654 assert_eq!(captured_f32(&received, 0), vec![0.5]);
655 let received = received.lock().unwrap();
656 let MediaBuffer::Audio(output) = &received[0] else {
657 panic!("expected audio")
658 };
659 assert!(!Arc::ptr_eq(&original, output));
660 }
661
662 #[test]
663 fn integer_amplification_saturates_instead_of_wrapping() {
664 let options = AudioVolumeOptions {
665 gain: 2.0,
666 ramp_duration: Duration::ZERO,
667 ..AudioVolumeOptions::default()
668 };
669 let (mut volume, _, received) = new_volume(options);
670 let mut frame = ffmpeg::frame::Audio::new(
671 ffmpeg::format::Sample::I16(Type::Packed),
672 2,
673 ffmpeg::ChannelLayout::MONO,
674 );
675 frame.set_rate(48_000);
676 for (destination, value) in frame.data_mut(0)[..4]
677 .chunks_exact_mut(2)
678 .zip([20_000_i16, -20_000_i16])
679 {
680 destination.copy_from_slice(&value.to_ne_bytes());
681 }
682 volume.consume(MediaBuffer::Audio(Arc::new(frame))).unwrap();
683
684 let received = received.lock().unwrap();
685 let MediaBuffer::Audio(frame) = &received[0] else {
686 panic!("expected audio")
687 };
688 let output: Vec<_> = frame.data(0)[..4]
689 .chunks_exact(2)
690 .map(|bytes| i16::from_ne_bytes(bytes.try_into().unwrap()))
691 .collect();
692 assert_eq!(output, vec![i16::MAX, i16::MIN]);
693 }
694
695 #[test]
696 fn rejects_invalid_controls_and_non_audio_buffers() {
697 let (_, handle) = AudioVolume::new("volume");
698 assert_eq!(
699 handle.set_gain(-1.0),
700 Err(AudioVolumeError::InvalidGain(-1.0))
701 );
702 assert!(matches!(
703 handle.set_gain_db(f32::NAN),
704 Err(AudioVolumeError::InvalidGainDb(value)) if value.is_nan()
705 ));
706
707 let (mut volume, _) = AudioVolume::new("volume");
708 let error = volume
709 .consume(MediaBuffer::Packet(Arc::new(ffmpeg::Packet::empty())))
710 .unwrap_err();
711 assert!(matches!(
712 error,
713 crate::Error::AudioVolumeError(AudioVolumeError::UnsupportedBuffer("Packet"))
714 ));
715 }
716}